import numpy as np
import matplotlib.pyplot as plt
from ruamel.yaml import YAML
from BeaverFramework.Framework.model import Model
from BeaverFramework.Framework.assembly import Assembly
from BeaverFramework.Framework.framework import Framework
from BeaverFramework.Framework.utles import load_config
from os.path import join
import os
from dotenv import load_dotenv
"""
authors @ MicahelZhang
"""
class MassSpringDamperModel(Model):
    def __init__(self, name, config):
        super().__init__(name)
        self.config = config

    def setup(self):
        self.mass = self.config['mass']
        self.damping = self.config['damping']
        self.spring_constant = self.config['spring_constant']
        self.inputs['initial_position'] = self.config['initial_position']
        self.inputs['initial_velocity'] = self.config['initial_velocity']
        self.inputs['t_span'] = tuple(self.config['t_span'])
        self.inputs['dt'] = self.config['dt']
        self.outputs['time'] = None
        self.outputs['results'] = None

    def equations_of_motion(self, t, y):
        x, v = y
        dxdt = v
        dvdt = (-self.damping * v - self.spring_constant * x) / self.mass
        return np.array([dxdt, dvdt])

    def runge_kutta(self, t, y, dt):
        k1 = self.equations_of_motion(t, y)
        k2 = self.equations_of_motion(t + dt/2, y + dt/2 * k1)
        k3 = self.equations_of_motion(t + dt/2, y + dt/2 * k2)
        k4 = self.equations_of_motion(t + dt, y + dt * k3)
        return y + (dt / 6) * (k1 + 2*k2 + 2*k3 + k4)

    def compute(self):
        t_span = self.inputs['t_span']
        dt = self.inputs['dt']
        y0 = np.array([self.inputs['initial_position'], self.inputs['initial_velocity']])

        time = np.arange(t_span[0], t_span[1], dt)
        results = np.zeros((len(time), len(y0)))

        y = y0
        for i, t in enumerate(time):
            results[i] = y
            y = self.runge_kutta(t, y, dt)

        self.outputs['time'] = time
        self.outputs['results'] = results

class PostProcessingModel(Model):
    def __init__(self, name, config):
        super().__init__(name)
        self.config = config

    def setup(self):
        self.inputs['time'] = None
        self.inputs['results'] = None
        self.figure_size = tuple(self.config['figure_size'])

    def compute(self):
        time = self.inputs['time']
        results = self.inputs['results']

        if time is None or results is None:
            print("Error: Missing data for plotting")
            return

        plt.figure(figsize=self.figure_size)
        plt.subplot(2, 1, 1)
        plt.plot(time, results[:, 0], label='Position (x)')
        plt.title('Mass-Spring-Damper System Response')
        plt.xlabel('Time (s)')
        plt.ylabel('Position (m)')
        plt.grid()
        plt.legend()

        plt.subplot(2, 1, 2)
        plt.plot(time, results[:, 1], label='Velocity (v)', color='orange')
        plt.xlabel('Time (s)')
        plt.ylabel('Velocity (m/s)')
        plt.grid()
        plt.legend()

        plt.tight_layout()


